/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lvyh.lightframe.tracer.plugin.springmvc;


import com.lvyh.lightframe.tracer.core.SimpleSpan;
import com.lvyh.lightframe.tracer.core.SimpleSpanContext;
import com.lvyh.lightframe.tracer.core.SimpleTracer;
import com.lvyh.lightframe.tracer.core.common.AbstractTracer;
import com.lvyh.lightframe.tracer.core.common.SimpleTracerConfiguration;
import com.lvyh.lightframe.tracer.core.common.TraceContants;
import io.opentracing.propagation.Format;

import javax.servlet.Filter;
import javax.servlet.FilterChain;
import javax.servlet.FilterConfig;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.IOException;
import java.util.Enumeration;
import java.util.HashMap;

public class SpringMvcTracerFilter implements Filter {

    private AbstractTracer springMvcTracer;

    public SpringMvcTracerFilter(AbstractTracer springMvcTracer) {
        this.springMvcTracer = springMvcTracer;
    }

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {
    }

    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain) {

        SimpleSpan springMvcSpan = null;
        long responseSize = -1;
        int httpStatus = -1;
        try {
            HttpServletRequest request = (HttpServletRequest) servletRequest;
            HttpServletResponse response = (HttpServletResponse) servletResponse;
            SimpleSpanContext spanContext = extractSpanContext(request);
            springMvcSpan = springMvcTracer.serverReceive(spanContext);

            springMvcSpan.setOperationName(request.getRequestURL().toString());
            springMvcSpan.setTag(TraceContants.LOCAL_APP, SimpleTracerConfiguration.getApplicationName());
            springMvcSpan.setTag(TraceContants.REQUEST_URL, request.getRequestURL().toString());
            springMvcSpan.setTag(TraceContants.METHOD, request.getMethod());
            springMvcSpan.setTag(TraceContants.REQ_SIZE, request.getContentLength());

            ResponseWrapper responseWrapper = new ResponseWrapper(response);

            filterChain.doFilter(servletRequest, response);

            httpStatus = responseWrapper.getStatus();
            responseSize = responseWrapper.getContentLength();
        } catch (Throwable t) {
            httpStatus = 500;
            throw new RuntimeException(t);
        } finally {
            if (springMvcSpan != null) {
                springMvcSpan.setTag(TraceContants.RESP_SIZE, responseSize);
                //ss
                springMvcTracer.serverSend(String.valueOf(httpStatus));
            }
        }
    }

    @Override
    public void destroy() {
    }


    public String getFilterName() {
        return "SpringMvcTracerFilter";
    }

    /***
     * Extract tracing context from request received from previous node
     */
    public SimpleSpanContext extractSpanContext(HttpServletRequest request) {
        HashMap<String, String> headers = new HashMap<String, String>();
        Enumeration headerNames = request.getHeaderNames();
        while (headerNames.hasMoreElements()) {
            String key = (String) headerNames.nextElement();
            String value = request.getHeader(key);
            headers.put(key, value);
        }

        if (headers.isEmpty()) {
            return null;
        }

        SimpleTracer tracer = springMvcTracer.getSimpleTracer();
        SimpleSpanContext spanContext = (SimpleSpanContext) tracer.extract(
                Format.Builtin.HTTP_HEADERS, new SpringMvcHeadersCarrier(headers));
        return spanContext;
    }


    class ResponseWrapper extends HttpServletResponseWrapper {

        int contentLength = 0;

        public ResponseWrapper(HttpServletResponse httpServletResponse) throws IOException {
            super(httpServletResponse);
        }

        @Override
        public void setContentLength(int len) {
            contentLength = len;
            super.setContentLength(len);
        }

        public int getContentLength() {
            return contentLength;
        }
    }

}